from sqlalchemy import create_engine, text, inspect
from sqlalchemy.orm import sessionmaker

from db.tables import Base

SQLALCHEMY_DATABASE_URL = "mysql+mysqlconnector://root:123456@localhost/ddos"
engine = create_engine(SQLALCHEMY_DATABASE_URL)
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)


def get_db():
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()


def init_database():
    # 创建连接到 MySQL 的引擎
    eng = create_engine("mysql+mysqlconnector://root:123456@localhost/mysql")
    with eng.connect() as connection:
        # 创建数据库，如果不存在
        connection.execute(text("CREATE DATABASE IF NOT EXISTS ddos CHARACTER SET utf8mb4 COLLATE utf8mb4_bin;"))
    eng.dispose()  # 关闭初始数据库连接

    # 连接到新创建的数据库
    eng = create_engine("mysql+mysqlconnector://root:123456@localhost/ddos")

    # 检查数据库中已存在的表并创建不存在的表
    inspector = inspect(eng)
    existing_tables = inspector.get_table_names()
    for table_name, table_obj in Base.metadata.tables.items():
        if table_name not in existing_tables:
            table_obj.create(bind=eng)


def init_sql_file():
    pass


def get_session():
    Session = sessionmaker(bind=engine)
    return Session()


if __name__ == '__main__':
    init_database()
